import csv
import numpy as np

def write_header(file_path):
    """
    将指定的 header 写入到文件头部.

    参数:
    file_path (str): 输出文件的路径.
    header (str): 要写入的 header 内容.
    """
    header="""/*--------------------------------*- C++ -*----------------------------------*\\
| =========                 |                                                 |
| \\      /  F ield         | OpenFOAM: The Open Source CFD Toolbox           |
|  \\    /   O peration     | Version:  v2306                                 |
|   \\  /    A nd           | Website:  www.openfoam.com                      |
|    \\/     M anipulation  |                                                 |
\\*---------------------------------------------------------------------------*/
FoamFile
{
    version     2.0;
    format      ascii;
    class       dictionary;
    object      fvOptions;
}
// * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * //

"""

    with open(file_path, 'w') as file:
        file.write(header)

def read_csv_column(file_path, column_index):
    """
    从 CSV 文件中读取指定列的数据到 NumPy 数组.

    参数:
    file_path (str): CSV 文件的路径.
    column_index (int): 要读取的列索引(从 0 开始计数).

    返回:
    numpy.ndarray: 指定列的数据.
    """
    with open(file_path, 'r') as file:
        reader = csv.reader(file)
        column_data = []
        for row in reader:
            column_data.append(row[column_index])

    return np.array(column_data)

def write_csv_file(file_path, x, y, z):
    """
    将 x, y, z 数组组合成 (x, y, z) 并写入到一个新的 CSV 文件中.

    参数:
    file_path (str): 输出 CSV 文件的路径.
    x (numpy.ndarray): x 坐标数据.
    y (numpy.ndarray): y 坐标数据.
    z (numpy.ndarray): z 坐标数据.
    """
    with open(file_path, 'a', newline='') as file:
        writer = csv.writer(file)
        # for row_x, row_y, row_z in zip(x, y, z):
        #     writer.writerow([f"({row_x} {row_y} {row_z})"])
        for i, ( row_x, row_y, row_z) in enumerate(zip(x[1:], y[1:], z[1:]), start=1):
            writer.writerow([f"location{i} ({row_x} {row_y} {row_z})"])
        writer.writerow([])

def write_injector(file_path, x):

    """
    将注入器位置写入到文件中.

    参数:
    file_path (str): 输出文件的路径.
    x (numpy.ndarray): 注入器位置的 x 坐标数据.
    """
    multi_lines_1="""{
    selectionMode   points;
    points
    (   
"""
    multi_lines_2="""    );
}
"""

    with open(file_path, 'a', newline='') as file:
        writer = csv.writer(file)
        for i, row_x in enumerate(x[1:], start=1):
            # print(row_x)
            writer.writerow([f"injector{i}"])
            file.write(multi_lines_1)
            writer.writerow([f"        $location{i}"])
            file.write(multi_lines_2)

def write_options(file_path,x,voc):

    # NOTE: Scaling the exp log, divided by scale_factor
    scale_factor = 100;

    multi_lines_1="""options
{
"""
    multi_lines_2="""    {
        type            scalarSemiImplicitSource;
"""
    multi_lines_3="""        volumeMode      absolute;
        sources
        {
"""
    multi_lines_4="""
        }
    }
"""
    multi_lines_last='}//End of options'

    #第一行的string不能转为float, 后面遍历的时候也不用跳过第一行
    new_voc = voc[1:].astype(float)
    with open(file_path, 'a', newline='') as file:
        writer = csv.writer(file)
        file.write(multi_lines_1)
        for i, (row_x, row_voc) in enumerate(zip(x[1:],new_voc), start=1):
            row_voc_change=row_voc/scale_factor
            writer.writerow([f"    scalarSource{i}"])
            file.write(multi_lines_2)
            writer.writerow([f"        $injector{i};"])
            file.write(multi_lines_3)
            writer.writerow([f"            C    ({row_voc_change} 0);"])
            file.write(multi_lines_4)
        file.write(multi_lines_last)

if __name__ == "__main__":
    input_file_path = "data.csv"
    output_file_path = "output.txt"
    column_index = [1, 2, 3]  # [x, y, z]
    x = read_csv_column(input_file_path, column_index[0])
    y = read_csv_column(input_file_path, column_index[1])
    z = read_csv_column(input_file_path, column_index[2])

    # 污染物
    voc = read_csv_column(input_file_path, 5)
    print(voc)

    write_header(output_file_path)
    write_csv_file(output_file_path, x, y, z)
    write_injector(output_file_path, x)
    write_options(output_file_path, x,voc)

